from datetime import datetime
import backtrader as bt
import matplotlib.pyplot as plt
import akshare as ak
import pandas as pd

from StdDevStrategy import StdDevStrategy
from data_feed import StockDataFetcher

plt.rcParams["font.sans-serif"] = ["SimHei"]
plt.rcParams["axes.unicode_minus"] = False
# 抓取股票
stock_fetcher = StockDataFetcher(symbol="000001")
data_stock = stock_fetcher.get_data()
#===============================================================================

# 初始化回测系统
cerebro = bt.Cerebro()

# 设置回测时间范围
start_date = datetime(2018, 4, 3)
end_date = datetime(2024, 5, 22)

# 加载数据到Cerebro
data = bt.feeds.PandasData(dataname=data_stock, 
                           fromdate=start_date, 
                           todate=end_date)
cerebro.adddata(data, name='000001')

# 加载策略到Cerebro
cerebro.addstrategy(StdDevStrategy)

# 设置初始资金和手续费
start_cash = 1000000 
cerebro.broker.setcash(start_cash)
cerebro.broker.setcommission(commission=0.002)

# 运行回测
results = cerebro.run()
#=====================================================================================
# 获取回测结果并打印
port_value = cerebro.broker.getvalue()
pnl = port_value - start_cash
print(f"初始资金: {start_cash}\n回测期间：{start_date.strftime('%Y%m%d')}:{end_date.strftime('%Y%m%d')}")
print(f"总资金: {round(port_value, 2)}")
print(f"净收益: {round(pnl, 2)}")
#============================================================
# 获取策略实例
strategy_instance = results[0]

# 调用策略实例的get_print_list方法获取print_list
# 调出成交列表
print_list = strategy_instance.get_print_list()

def parse_data_strings_to_df(data_strings):  # 成交列表转成dataframe
    data_rows = []  
    for data_str in data_strings:  
        parts = data_str.split(',')  # 按逗号分割字符串  
        row_data = {}  
        for part in parts:  
            key_value = part.split(':')  # 尝试按冒号分割键值对  
            if len(key_value) == 2:  
                key, value = key_value  
                row_data[key.strip()] = value.strip()  # 去除键和值两侧的空白字符  
            else:  
                print(f"Warning: Invalid key-value pair: {part}")  
        data_rows.append(row_data)  
    return pd.DataFrame(data_rows)  
  
# 创建一个DataFrame  
df = parse_data_strings_to_df(print_list) 
df.to_csv('print_list.csv', index=False)
# 打印DataFrame查看结果  
print(df)
#==================================================================

cerebro.plot(style='candlestick',figsize=(12, 8))
